import tensorflow as tf
from tensorflow.keras import layers, models, optimizers

class lenet(models.Model):
    def __init__(self, input_shape):
        super (lenet, self).__init__()
        self.c1 = layers.Conv2D(6, (5, 5), padding='same', activation='relu')
        self.s2 = layers.MaxPooling2D(pool_size=(2, 2), strides=2)
        self.c3 = layers.Conv2D(16, (5, 5), padding='same', activation='relu')
        self.s4 = layers.MaxPooling2D(pool_size=(2, 2), strides=2)
        self.falt = layers.Flatten()
        self.fc5 = layers.Dense(120, activation='relu')
        self.fc6 = layers.Dense(84, activation='relu')
        self.fc7 = layers.Dense(10, activation='softmax')

    def call(self, x, training=None, mask=None):
        out = self.c1(x)
        out = self.s2(out)
        out = self.c3(out)
        out = self.s4(out)
        out = self.falt(out)
        out = self.fc5(out)
        out = self.fc6(out)
        out = self.fc7(out)
        return out

from tensorflow.keras.datasets import mnist
from tensorflow.keras import utils

# 批处理量
batch_size = 128
nb_classes = 10
# 循环次数
nb_epoch = 12

# 图片维度设定
img_rows, img_cols = 28, 28
# 卷积核数量
nb_filters = 32
# 池化尺寸
pool_size = (2, 2)
# 卷积核尺寸
kernel_size = (3, 3)

# 切分数据集
(X_train, y_train), (X_test, y_test) = mnist.load_data()

X_train = X_train.astype('float32')
X_test = X_test.astype('float32')
# 归一化
X_train /= 255
X_test /= 255
# 将图像转换为  数量，宽，高，通道
X_train = X_train.reshape(-1, 28, 28, 1)
X_test = X_test.reshape(-1, 28, 28, 1)
print('X_train shape:', X_train.shape)
print(X_train.shape[0], 'train samples')
print(X_test.shape[0], 'test samples')

# 独热编码处理
Y_train = utils.to_categorical(y_train, nb_classes)
Y_test = utils.to_categorical(y_test, nb_classes)

model = lenet([28, 28, 1])

model.compile(loss='categorical_crossentropy',
              optimizer=optimizers.Adam(0.001),
              metrics=['accuracy'])

model.fit(X_train, Y_train, batch_size=batch_size, epochs=nb_epoch,
          verbose=1, validation_data=(X_test, Y_test))
score = model.evaluate(X_test, Y_test, verbose=0)
print('Test score:', score[0]) # 代价
print('Test accuracy:', score[1]) # 准确率
